import numpy as np
import pandas as pd
from sklearn.metrics import r2_score
from typing import Dict, List, Tuple, NamedTuple

class BenchmarkResult(NamedTuple):
    """基准测试结果数据类"""
    mean_score: float  # 平均R2得分
    max_score: float  # 最高R2得分
    min_score: float  # 最低R2得分

def evaluate_model(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """评估模型性能，返回R2分数"""
    return r2_score(y_true, y_pred)

def run_single_evaluation(df: pd.DataFrame, feature: callable, random_state: int = 42) -> Tuple[float, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """运行单个预测方法并评估"""
    from utils.model_utils import train_model, predict
    
    X, y = feature(df)
    
    # 训练预测评估
    model, X_train, X_test, y_train, y_test = train_model(X, y, random_state=random_state)
    predictions = predict(model, X_test)
    score = evaluate_model(y_test, predictions)
    
    # 获取测试集日期并确保顺序一致
    dates = pd.to_datetime(df.iloc[:, 0])
    test_indices = np.arange(len(dates))[len(X_train):len(X_train)+len(X_test)]
    test_dates = dates.iloc[test_indices]
    
    # 按日期排序所有数据
    sort_idx = np.argsort(test_dates)
    test_dates = test_dates.iloc[sort_idx]
    predictions = predictions[sort_idx]
    y_test = y_test[sort_idx]
    y_test_full = df.iloc[test_indices, 1].values[sort_idx]
    
    return score, y_test, predictions, test_dates, y_test_full

def run_benchmark_evaluation(df: pd.DataFrame, feature: callable, random_states: List[int]) -> List[float]:
    """运行基准测试评估"""
    scores = []
    for state in random_states:
        score, _, _, _, _ = run_single_evaluation(df, feature, state)
        scores.append(score)
    return scores

def analyze_benchmark_results(results: Dict[str, List[float]], n_trials: int) -> Dict[str, BenchmarkResult]:
    """分析基准测试结果"""
    analyzed_results = {}
    for method_name, scores in results.items():
        analyzed_results[method_name] = BenchmarkResult(
            mean_score=np.mean(scores),
            max_score=np.max(scores),
            min_score=np.min(scores)
        )
    return analyzed_results

def select_best_method(analyzed_results: Dict[str, BenchmarkResult]) -> Tuple[str, BenchmarkResult]:
    """选择最佳方法"""
    # 先只考虑正分数的方法
    valid_results = {name: result for name, result in analyzed_results.items() 
                    if result.mean_score > 0}
    
    # R2最高的
    if valid_results:
        return max(valid_results.items(), key=lambda x: x[1].mean_score)
    
    # 如果所有方法都是负分，选择最接近0的
    print("\n警告: 所有方法的R2得分都小于0，选择损失最小的方法")
    return min(analyzed_results.items(), key=lambda x: abs(x[1].mean_score))

def print_benchmark_results(results: Dict[str, List[float]], n_trials: int) -> None:
    """打印基准测试结果"""
    analyzed_results = analyze_benchmark_results(results, n_trials)
    
    print(f"\nResults({n_trials} trials):")
    
    # 结果表格
    rows = []
    for method, result in analyzed_results.items():
        rows.append({
            '方法': method,
            '平均R2': f"{result.mean_score:.4f}",
            '最高R2': f"{result.max_score:.4f}",
            '最低R2': f"{result.min_score:.4f}"
        })
    
    # 按平均分排序
    df_results = pd.DataFrame(rows)
    df_results = df_results.sort_values('平均R2', ascending=False)
    print("\n", df_results.to_string(index=False))
    
    best_method, best_result = select_best_method(analyzed_results)  
    print(f"\n选择方法: {best_method}")
    print(f"- 平均R2得分: {best_result.mean_score:.4f}")
    print(f"- 最高R2得分: {best_result.max_score:.4f}")
    print(f"- 最低R2得分: {best_result.min_score:.4f}")